import spacy
from elasticsearch import Elasticsearch
import numpy as np

# 加载中文语言模型
nlp = spacy.load("zh_core_web_md")

# 连接到 Elasticsearch
es = Elasticsearch([{'host': 'es-sg-wzf25x0az0020g2ww.elasticsearch.aliyuncs.com', 'port': 9200,
                     # 'username': 'elastic', 'password': 'AocDidyf5Hom68EC',
                     'scheme': 'http'}], basic_auth=('elastic', 'AocDidyf5Hom68EC'))

# 示例文档
documents = [
    {"id": 1, "text": "机器学习是一门涉及人工智能的领域，它关注如何通过经验自动改善计算机系统的性能。"},
    {"id": 2, "text": "自然语言处理是人工智能的一个子领域，研究计算机如何理解和生成人类语言。"},
    {"id": 3, "text": "信息检索是从大量文档中检索相关信息的过程，常用于构建搜索引擎。"},
    {"id": 4, "text": "深度学习是机器学习的一个分支，通过神经网络模拟人类的学习过程。"},
]

# 将文本向量存储到 Elasticsearch
for doc in documents:
    # 处理文本并获取文本向量
    doc_vector = nlp(doc["text"]).vector
    # 存储文本向量到 Elasticsearch
    es.index(index='document_vectors', id=doc["id"], body={'vector': doc_vector.tolist()})

# 用户查询
query = "深度学习在人工智能中的应用"

# 处理查询文本并获取查询向量
query_vector = nlp(query).vector

# 在 Elasticsearch 中进行信息检索
search_results = es.search(index='document_vectors', body={
    "query": {
        "script_score": {
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, doc['vector']) + 1.0",
                "params": {"query_vector": query_vector.tolist()}
            }
        }
    }
})

# 提取检索结果
hits = search_results['hits']['hits']
for hit in hits:
    document_id = hit['_id']
    document_score = hit['_score']
    print(f"文档ID: {document_id}, 相似度分数: {document_score}")
